#ifndef MY_PARCSR_HPP
#define MY_PARCSR_HPP

#include <cstdio>
#include <cstdlib>
#include <cassert>
#include "mpi.h"
#include <vector>
#include "cuda_common.hpp"
#include <unordered_map>
#include <string>
#include <cmath>

/// @brief 并行CSR向量
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam calc_t 数据类型，可以为float或者double
template<typename idx_t, typename calc_t>
class MyParCSRVector
{
public:
    const MPI_Comm comm = MPI_COMM_NULL;// 通信域
    const idx_t glb_nrows;// 向量的全局行数
    const idx_t beg_row;// 本进程负责的起始行号
    const idx_t end_row;// 本进程负责的末尾行号（注意是开区间）
    calc_t * data = nullptr;// 本进程负责范围的数据
    bool use_buf = false;

    MyParCSRVector(MPI_Comm _comm, const idx_t _glb_nrows, const idx_t _beg, const idx_t _end, const double * dist_data):
        comm(_comm), glb_nrows(_glb_nrows), beg_row(_beg), end_row(_end)
    {
        const idx_t nrows = end_row - beg_row;
#ifdef USE_GPU
        cuda_malloc((void **)(&data), nrows * sizeof(calc_t));
        if (dist_data) {
            if constexpr (sizeof(calc_t) == sizeof(double)) {
                cuda_memcpy(data, dist_data, nrows, 0);
            } else {
                double * data_setup;
                cuda_malloc((void **)(&data_setup), nrows * sizeof(double));
                cuda_memcpy(data_setup, dist_data, nrows, 0);
                cuda_vec_copy(data, data_setup, nrows);
                cuda_free(data_setup);
            }
        } 
#else
        data = new calc_t [nrows];
        if (dist_data) {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++)
                data[i] = dist_data[i];
        } else {// 提前赋零以免出nan
            // int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
            // printf("P%d make zero %p\n", my_pid, &data);
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++)
                data[i] = 0.0;
        }
#endif
    }
    MyParCSRVector(const MyParCSRVector<idx_t, calc_t> & x, calc_t *buf = nullptr):
        comm(x.comm), glb_nrows(x.glb_nrows), beg_row(x.beg_row), end_row(x.end_row)
    {
        const idx_t nrows = end_row - beg_row;
#ifdef USE_GPU
        if (buf) {
            this->use_buf = true;
            data = buf;
        } else {
            cuda_malloc((void **)(&data), nrows * sizeof(calc_t));
        }
#else
        if (buf) {
            this->use_buf = true;
            data = buf;
        } else {
            data = new calc_t [nrows];
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++)
                data[i] = 0.0;
        }
#endif
    }
    ~MyParCSRVector() {
#ifdef USE_GPU
        if (!use_buf && data) {
            cuda_free(data);
        }
#else
        if (!use_buf && data) {
            delete [] data;
        } 
#endif
    }
    void set_val(const calc_t val) {
#ifdef USE_GPU
        assert(false);
#else
        assert(data);
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < end_row - beg_row; i++)
            data[i] = val;
#endif
    }
};



/// @brief 向量点积
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param x 向量一
/// @param y 向量二
/// @return 内积结果 (x,y)
template<typename idx_t, typename data_t>
data_t vec_dot(const MyParCSRVector<idx_t, data_t> & x, const MyParCSRVector<idx_t, data_t> & y)
{
    assert(x.glb_nrows == y.glb_nrows && x.beg_row == y.beg_row && x.end_row == y.end_row);
    const idx_t nrows = x.end_row - x.beg_row;
    data_t loc_res = 0.0, glb_res;
#ifdef USE_GPU
    loc_res = cuda_vec_dot<idx_t, data_t, data_t>(x.data, y.data, nrows);
    return loc_res;
#else
    #pragma omp parallel for schedule(static) reduction(+:loc_res)
    for (idx_t i = 0; i < nrows; i++)
        loc_res += (data_t) x.data[i] * (data_t) y.data[i];
    MPI_Allreduce(&loc_res, &glb_res, 1, (sizeof(data_t) == 4) ? MPI_FLOAT : MPI_DOUBLE, MPI_SUM, x.comm);
    return glb_res;
#endif
}

/// @brief 向量范数
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param x 向量
/// @return 向量范数 ||x||_2
template<typename idx_t, typename data_t>
double vec_norm(const MyParCSRVector<idx_t, data_t> & x)
{
    return sqrt(vec_dot(x, x));
}

/// @brief 向量加法
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param v1 源向量一
/// @param alpha 系数
/// @param v2 源向量二
/// @param v 目标向量
/// @result v = v1 + alpha * v2 
template<typename idx_t, typename data_t>
void vec_add(const MyParCSRVector<idx_t, data_t> & v1, data_t alpha,
            const MyParCSRVector<idx_t, data_t> & v2,
            MyParCSRVector<idx_t, data_t> & v)
{
    const idx_t nrows = v.end_row - v.beg_row;

#ifdef USE_GPU
    cuda_vec_add(v.data, v1.data, v2.data, alpha, nrows);
#else
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows; i++)
        v.data[i] = v1.data[i] + alpha * v2.data[i];
#endif
}

/// @brief 向量加法
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param x 源向量
/// @param alpha 系数
/// @param y 目标向量
/// @result y = alpha * x
template<typename idx_t, typename data_t>
void vec_scale(const MyParCSRVector<idx_t, data_t> & x, data_t alpha,
            MyParCSRVector<idx_t, data_t> & y)
{
    const idx_t nrows = y.end_row - y.beg_row;

#ifdef USE_GPU
    cuda_vec_scale(y.data, x.data, alpha, nrows);
#else
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows; i++)
        y.data[i] = alpha * x.data[i];
#endif
}

/// @brief 向量拷贝（从源向量到目标向量的拷贝）
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t1 源向量的数据类型，可以为float或者double
/// @tparam data_t2 目标向量的数据类型，可以为float或者double
/// @param src 源向量
/// @param dst 目标向量
template<typename idx_t, typename data_t1, typename data_t2>
void vec_copy(const MyParCSRVector<idx_t, data_t1> & src, MyParCSRVector<idx_t, data_t2> & dst)
{
    assert(src.glb_nrows == dst.glb_nrows && src.beg_row == dst.beg_row && src.end_row == dst.end_row);
    const idx_t nrows = dst.end_row - dst.beg_row;

#ifdef USE_GPU
    cuda_vec_copy(dst.data, src.data, nrows);
#else
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows; i++)
        dst.data[i] = src.data[i];
#endif
}

/// @brief 本地CSR矩阵
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 数据存储类型，可以为float或者double
/// @tparam calc_t 数据计算类型，可以为float或者double
template<typename idx_t, typename data_t, typename calc_t>
class MyCSRMatrix
{
public:
    const idx_t nrows;// 行数
    idx_t ncols;// 列数
    idx_t nnz;// 非零元数
    idx_t * row_ptr = nullptr;// 行偏移量
    idx_t * col_idx = nullptr;// 列索引
    data_t* vals = nullptr;// 非零元值

    /// @brief 本地CSR矩阵的构造函数，只开辟内存，不做初始化
    /// @param _nrows 行数
    /// @param _ncols 列数
    /// @param _nnz 非零元数
    MyCSRMatrix(const idx_t _nrows, const idx_t _ncols, const idx_t _nnz): nrows(_nrows), ncols(_ncols), nnz(_nnz) {
#ifdef USE_GPU
        cuda_malloc((void **)(&row_ptr), (nrows + 1) * sizeof(idx_t));
        cuda_malloc((void **)(&col_idx), nnz * sizeof(idx_t));
        cuda_malloc((void **)(&vals), nnz * sizeof(data_t));
#else
        row_ptr = new idx_t [nrows + 1];
        col_idx = new idx_t [nnz];
        vals    = new data_t[nnz];
#endif
    }
    MyCSRMatrix(const idx_t _nrows, const idx_t _ncols): nrows(_nrows), ncols(_ncols) {
        row_ptr = new idx_t [nrows + 1];
    }
    // 通过转置生成（没有数值）的CSR矩阵，map长度应为len
    MyCSRMatrix(const idx_t _nrows, const idx_t len, const idx_t * map_f2c):
        nrows(_nrows), ncols(len), nnz(len) {
        assert(map_f2c);

        row_ptr = new idx_t [nrows + 1]; row_ptr[0] = 0;
        idx_t * cnt = new idx_t [nrows];
        for (idx_t i = 0; i < nrows; i++) {// 初始化清零
            row_ptr[i + 1] = 0;
            cnt[i] = 0;
        }
        // 累加计数
        for (idx_t j = 0; j < len; j++) {
            idx_t rid = map_f2c[j];// 目标行号
            assert(rid < nrows);
            row_ptr[rid + 1] ++;
        }
        // 累加行偏移
        for (idx_t i = 0; i < nrows; i++) {
            row_ptr[i + 1] += row_ptr[i];
            cnt[i] = 0;
        }
        // 整理列序号
        col_idx = new idx_t [nnz];
        for (idx_t j = 0; j < len; j++) {
            idx_t rid = map_f2c[j];// 目标行号
            col_idx[ row_ptr[rid] + cnt[rid] ] = j;
            cnt[rid] ++;
        }
#ifdef DEBUG_PRINT
        for (idx_t i = 0; i < nrows; i++) {
            assert(cnt[i] == row_ptr[i + 1] - row_ptr[i]);
            // if (row_ptr[i] >= row_ptr[i + 1]) {
            //     int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, & my_pid);
            //     printf("Proc %d Error i %d : %d %d\n", my_pid, i, row_ptr[i], row_ptr[i + 1]);
            // }
            assert(row_ptr[i] < row_ptr[i + 1]);
        }
#endif
        delete [] cnt;
    }
    ~MyCSRMatrix() {
#ifdef USE_GPU
        if (row_ptr) cuda_free(row_ptr);
        if (col_idx) cuda_free(col_idx);
        if (vals)    cuda_free(vals);
#else
        if (row_ptr) delete [] row_ptr;
        if (col_idx) delete [] col_idx;
        if (vals)    delete [] vals;
#endif
    }
    // y := alpha*A*x + beta*b
    void Mult(calc_t alpha, const calc_t *x, calc_t beta, const calc_t * b, calc_t *y) {
#ifdef USE_GPU
        cuda_spmv(row_ptr, col_idx, vals, x, b, y, alpha, beta, nrows);
#else
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < nrows; i++) {
            calc_t tmp = 0.0;
            __builtin_prefetch(col_idx + row_ptr[i + 1], 0, 0);
            __builtin_prefetch(vals    + row_ptr[i + 1], 0, 0);
            for (idx_t p = row_ptr[i]; p < row_ptr[i+1]; p++) {
                tmp += vals[p] * x[col_idx[p]];
            }
            y[i] = alpha * tmp + beta * b[i];
        }
#endif
    }
    // y := alpha*A*x + beta*y
    void Mult(calc_t alpha, const calc_t * x, calc_t beta, calc_t *y) {
        Mult(alpha, x, beta, y, y);
    }
};

/// @brief 三维点
/// @tparam data_t 坐标数据类型 
template<typename data_t>
struct Point {
    data_t coord[3]; // 三维坐标
    /// @brief 三维点的构造函数
    Point() { coord[0] = coord[1] = coord[2] = 0.0; }
    /// @brief 三维点的构造函数
    /// @param _x, _y, _z 三维坐标
    Point(data_t _x, data_t _y, data_t _z): coord {_x, _y, _z} {}
    /// @brief 获取与点 b 的距离
    /// @param b 点 b
    data_t dist(const Point & b) const {
        const data_t diff[3] = {coord[0] - b.coord[0],
                                coord[1] - b.coord[1],
                                coord[2] - b.coord[2] };
        data_t t = diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2];
        // return sqrt(t);
        return t;
    }
    /// @brief 获取与点 b 的中点
    /// @param b 点 b
    Point middle(const Point & b) const {
        Point ret;
        for (int i = 0; i < 3; i++) ret.coord[i] = (coord[i] + b.coord[i]) / 2;
        return ret;
    }
    /// @brief 重载等号操作符
    /// @param b 点 b
    bool operator == (const Point & b) const {
        return coord[0] == b.coord[0] && coord[1] == b.coord[1] && coord[2] == b.coord[2];
    }
    /// @brief 加法操作 自身 += b
    /// @param b 点 b
    void add(const Point & b) {
        coord[0] += b.coord[0];
        coord[1] += b.coord[1];
        coord[2] += b.coord[2];
    }
    /// @brief 除法操作 自身 /= num
    /// @param num 除数
    Point div(const data_t num) {
        Point ret;
        ret.coord[0] = coord[0] / num;
        ret.coord[1] = coord[1] / num;
        ret.coord[2] = coord[2] / num;
        return ret;
    }
};

/// @brief 并行CSR矩阵
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 数据存储类型，可以为float或者double
/// @tparam calc_t 数据计算类型，可以为float或者double
template<typename idx_t, typename data_t, typename calc_t>
class MyParCSRMatrix
{
public:
    MPI_Comm comm = MPI_COMM_NULL;// 通信域
    idx_t glb_nrows;// 矩阵的全局行数
    idx_t beg_row;// 本进程负责的起始行号
    idx_t end_row;// 本进程负责的末尾行号（注意是开区间）
    idx_t * rows_partition = nullptr;// 记录各个进程的起始行号
    
    idx_t num_recvs = 0;// 有多少个邻居需要从他们那收取数据
    int   * recv_pids = nullptr;// 需要从哪些进程收？
    idx_t * recv_cnts = nullptr;// 从其它的每个进程处需要收多少个向量元素
    idx_t * col_map_offd = nullptr;// 从每个进程收的元素的行号，同时也记录了将全局坐标映射为非对角块格式的列序号

    idx_t num_sends = 0;// 有多少个邻居需要给他们发送自己负责的数据
    int   * send_pids = nullptr;// 需要给哪些进程发？
    idx_t * send_cnts = nullptr;// 给每个进程发的数量
    idx_t * send_row_ids = nullptr;// 给每个进程发的元素的行号

    MyCSRMatrix<idx_t, data_t, calc_t> * diag = nullptr;// 对角块矩阵
    MyCSRMatrix<idx_t, data_t, calc_t> * offd = nullptr;// 非对角块矩阵

    calc_t * send_buf = nullptr;// 发送缓冲区
    MPI_Request * send_reqs = nullptr;// 发送请求的句柄
    MPI_Request * recv_reqs = nullptr;// 接受请求的句柄

    Point<data_t> * points = nullptr;

    /// @brief 并行CSR矩阵的构造函数
    /// @param _comm 通信域
    /// @param _glb_nrows 矩阵的全局行数 （隐含列数==行数）
    /// @param _beg 本进程负责的起始行号
    /// @param _end 本进程负责的末尾行号（注意是开区间）
    /// @param dist_row_ptr 分布式的行偏移
    /// @param dist_col_idx 分布式的列索引
    /// @param dist_vals 分布式的非零元值
    MyParCSRMatrix(MPI_Comm _comm, const idx_t _glb_nrows, const idx_t _beg, const idx_t _end,
        const idx_t * dist_row_ptr, const idx_t * dist_col_idx, const double * dist_vals);
    MyParCSRMatrix(MPI_Comm _comm);
    ~MyParCSRMatrix();

    /// @brief 初始化本进程负责的行所对应的点的三维坐标（从文件读入）
    /// @param filename 文件名
    void init_point_coords(const std::string filename);

    /// @brief 建立通信所需的元数据
    /// @param gid2pidlid 记录了从全局序号到负责进程号和局部序号的映射（哈希表）
    void buildup_comm_handles(std::unordered_map<idx_t, std::pair<int, idx_t> > & gid2pidlid);

    /// @brief 将自己负责区域、而又需被别人引用到的元素发出去，同时接收自己需要引用到别的进程的数据
    /// @param my_vec 本进程负责范围的向量，长度应为end_row - beg_row
    /// @param other_vec 其它进程负责范围的而本进程需要引用到的向量，长度应为offd->ncols
    void update_halo(const calc_t * my_vec, calc_t * other_vec) const;
    
    /// @brief 启动halo通信（不保证完成）
    /// @param my_vec 本进程负责范围的向量，长度应为end_row - beg_row
    /// @param other_vec 其它进程负责范围的而本进程需要引用到的向量，长度应为offd->ncols
    void init_update_halo(const calc_t * my_vec, calc_t * other_vec) const;

    /// @brief 等待上一次调用init_update_halo通信的完成
    void wait_update_halo() const;

    // y := alpha*A*x + beta*b
    void Mult(calc_t alpha, const MyParCSRVector<idx_t, calc_t> & x,
            calc_t beta, const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & y) const;
    // y := alpha*A*x + beta*y
    void Mult(calc_t alpha, const MyParCSRVector<idx_t, calc_t> & x,
            calc_t beta, MyParCSRVector<idx_t, calc_t> & y) const;
};
typedef MyParCSRMatrix<int, float , double> MyParCSRMatrix_I32F32F64;
typedef MyParCSRMatrix<int, double, double> MyParCSRMatrix_I32F64F64;
typedef MyParCSRMatrix<int, float , float > MyParCSRMatrix_I32F32F32;

#endif // MY_PARCSR_HPP